package weka.classifiers.collective.meta;

import java.io.Serializable;
import java.util.Enumeration;
import java.util.List;
import java.util.Vector;

import weka.classifiers.Classifier;
import weka.classifiers.Evaluation;
import weka.classifiers.IOutOfBagCalculator;
import weka.classifiers.collective.*;

import weka.core.AdditionalMeasureProducer;
import weka.core.DiversityMeasureImplementation;
import weka.core.Instance;
import weka.core.Instances;
import weka.core.Option;
import weka.core.Utils;

public class ValidationSetEvaluator extends CollectiveRandomizableSingleClassifierEnhancer implements
        AdditionalMeasureProducer {

    /**
     *
     */
    private static final long serialVersionUID = 3466135436137638084L;
    private double m_validationPercent;
    private boolean[][] m_inBag;
    private SetEvaluator trainSet, validationSet, outOfBagSet, testSet,
            trainClassifierSet, validationClassifierSet, testClassifierSet;

    private class SetEvaluator implements Serializable {

        private Instances set;
        private DiversityMeasureImplementation diversity;
        private double accuracy, baseAccuracy, brier, brierdiversity;
        private double[] auc;

        public DiversityMeasureImplementation getDiversity() {
            return diversity;
        }

        public double getBrierScore() {
            return brier;
        }

        public double getBrierDiversity() {
            return brierdiversity;
        }

        public double getBaseAccuracy() {
            return baseAccuracy;
        }

        public double[] getAUC() {
            return auc;
        }

        public double getAccuracy() {
            return accuracy;
        }

        public SetEvaluator(Instances set) throws Exception {
            this.set = set;

            initializeDiversityMeasureImplementation(m_Classifier, set);
            diversity = getDiversityMeasureImplementation();

            accuracy = calcAccuracy(m_Classifier, set);
            for (Classifier base : m_Classifier.getEnsembleMembers()) {
                baseAccuracy += calcAccuracy(base, set);
            }
            baseAccuracy /= m_Classifier.getEnsembleMembers().size();
            auc = new double[set.numClasses()];
            brier = calcBrierScore();
            brierdiversity = calcBrierDiversity();
            //for (int classIndex = 0; classIndex < set.numClasses(); classIndex++)
            //	auc[classIndex] = eval.areaUnderROC(classIndex);
        }

        public SetEvaluator(Instances set, boolean[][] inBag) throws Exception {
            this.set = set;

            initializeDiversityMeasureImplementation(m_Classifier, set, inBag);
            diversity = getDiversityMeasureImplementation();
            accuracy = 100 * (1 - ((IOutOfBagCalculator) m_Classifier).measureOutOfBagError());
            List<Classifier> classifiers = m_Classifier.getEnsembleMembers();
            for (int i = 0; i < classifiers.size(); i++) {
                Classifier base = classifiers.get(i);
                Instances outOfBag = set.stringFreeStructure();
                for (int j = 0; j < set.numInstances(); j++) {
                    if (!inBag[i][j]) {
                        outOfBag.add(set.get(j));
                    }
                }
                baseAccuracy += calcAccuracy(base, outOfBag);
            }
            baseAccuracy /= m_Classifier.getEnsembleMembers().size();
            auc = new double[set.numClasses()];
        }

        private double calcAccuracy(Classifier classifier, Instances set) throws Exception {
            if (set.isEmpty()) {
                return 0;
            }
            double accuracy = 0;
            for (int i = 0; i < set.numInstances(); i++) {
                if (classifier.classifyInstance(set.get(i)) == set.get(i).classValue()) {
                    accuracy++;
                }
            }
            return (accuracy * 100) / set.numInstances();
        }

//    varje ensemblemedlem ger ifrån sig vektorn p_i = <p_i1, ..., p_ik>
//    (dvs. en sannolikhetsfördelning över de k klasserna, och i fallet med
//            klassificeringsträd så är en av sannolikheterna p_ij = 1 och övriga 0),
//            och medelvektorn p_m = <p_m1, ..., p_mk>, och klassvektorn är c = <c_1, ..., c_k>
//            (där den korrekta klassen har värdet 1 och övriga 0). Diversiteten blir
//            ((p_m-p_1)^2 + ... + (p_m-p_n)^2)/n (där n är antalet ensemblemedlemmar)
//            och Brier score blir (p_m-c)^2.

        // Calculation based on http://en.wikipedia.org/wiki/Brier_score
        public double calcBrierScore() throws Exception {
            double[] ensPred;
            double briersum = 0;
            for (int i = 0; i < set.numInstances(); i++) {
                ensPred = m_Classifier.distributionForInstance(set.instance(i));
                double briertmp = 0;
                for (int c = 0; c < set.numClasses(); c++) {
                    if (set.instance(i).classValue() == c) {
                        briertmp += Math.pow(ensPred[c] - 1, 2);
                    } else {
                        briertmp += Math.pow(ensPred[c] - 0, 2);
                    }
                }
                briersum += briertmp;
            }
            return briersum/set.numInstances();
        }

        public double calcBrierDiversity() throws Exception {
            List<Classifier> classifiers = m_Classifier.getEnsembleMembers();
            double[] basePred, ensPred;
            double brierdiv = 0;
            
            for (int i = 0; i < set.numInstances(); i++) {
                ensPred = m_Classifier.distributionForInstance(set.instance(i));
                double briertmp = 0;
                for (int b = 0; b < classifiers.size(); b++) {
                    basePred = classifiers.get(b).distributionForInstance(set.instance(i));
                    for (int c = 0; c < set.numClasses(); c++) {                        
                        briertmp += Math.pow(ensPred[c] - basePred[c], 2);
                    }
                }
                brierdiv += briertmp/classifiers.size();
            }
            return brierdiv/set.numInstances();
        }
    }

    @Override
    protected void build() throws Exception {
        buildClassifier();
    }

    @Override
    protected void buildClassifier() throws Exception {
        setMeasureDiversity(true);
        Instances train, validation;
        if (!m_Classifier.isEnsemble() & getMeasureDiversity()) {
            throw new Exception("Classifier to be trained need to be an ensemble.");
        }
        if ((m_validationPercent < 0) || (m_validationPercent > 100)) {
            throw new Exception("Percentage must be between 0 and 100");
        } else if (m_validationPercent > 0) {
            int trainSize = (int) Math.round(m_Trainset.numInstances()
                    * (100 - m_validationPercent) / 100);
            train = m_Trainset.stringFreeStructure();
            validation = m_Trainset.stringFreeStructure();
            for (int i = 0; i < trainSize; i++) {
                train.add(m_Trainset.get(i));
            }
            for (int i = trainSize; i < m_Trainset.numInstances(); i++) {
                validation.add(m_Trainset.get(i));
            }
        } else {
            train = m_Trainset;
            validation = m_Trainset.stringFreeStructure();
        }

        m_Classifier.buildClassifier(train);

        trainSet = new SetEvaluator(train);
        validationSet = new SetEvaluator(validation);
        testSet = new SetEvaluator(m_TestsetOriginal);

        trainClassifierSet = new SetEvaluator(setClassifierTarget(new Instances(train)));
        validationClassifierSet = new SetEvaluator(setClassifierTarget(new Instances(validation)));
        testClassifierSet = new SetEvaluator(setClassifierTarget(new Instances(m_TestsetOriginal)));
        if (m_Classifier instanceof IOutOfBagCalculator) {
            m_inBag = ((IOutOfBagCalculator) m_Classifier).getOutOfBag();
            outOfBagSet = new SetEvaluator(train, m_inBag);
        }
    }

    private Instances setClassifierTarget(Instances instances) throws Exception {
        for (int i = 0; i < instances.numInstances(); i++) {
            Instance inst = instances.get(i);
            inst.setClassValue(classifyInstance(inst));
        }
        return instances;
    }

    private DiversityMeasureImplementation evaluateOnSet(Instances set,
            Double acc) throws Exception {
        this.initializeDiversityMeasureImplementation(m_Classifier, set);

        Evaluation eval = new Evaluation(set);
        eval.setPriors(m_Trainset);
        eval.evaluateModel(m_Classifier, set);
        acc = new Double(eval.pctCorrect());

        return this.getDiversityMeasureImplementation();
    }

    private DiversityMeasureImplementation evaluateOnSet(Instances m_Trainset,
            boolean[][] m_inBag) {
        return null;
    }

    @Override
    protected double[] getDistribution(Instance instance) throws Exception {
        return m_Classifier.distributionForInstance(instance);
    }

    @Override
    public Enumeration enumerateMeasures() {
        m_AdditionalMeasures.add("measureAccuracyTrain");
        m_AdditionalMeasures.add("measureMeanAccuracyTrain");
        m_AdditionalMeasures.add("measureAccuracyDiffTrain");
        m_AdditionalMeasures.add("measureAUCTrain");
        m_AdditionalMeasures.add("measureBrierScoreTrain");
        m_AdditionalMeasures.add("measureBrierDiversityTrain");
        Enumeration<String> measures = new DiversityMeasureImplementation().enumerateMeasures();
        while (measures.hasMoreElements()) {
            m_AdditionalMeasures.add(measures.nextElement() + "Train");
        }

        m_AdditionalMeasures.add("measureAccuracyValidation");
        m_AdditionalMeasures.add("measureMeanAccuracyValidation");
        m_AdditionalMeasures.add("measureAccuracyDiffValidation");
        m_AdditionalMeasures.add("measureAUCValidation");
        m_AdditionalMeasures.add("measureBrierScoreValidation");
        m_AdditionalMeasures.add("measureBrierDiversityValidation");
        measures = new DiversityMeasureImplementation().enumerateMeasures();
        while (measures.hasMoreElements()) {
            m_AdditionalMeasures.add(measures.nextElement() + "Validation");
        }

        m_AdditionalMeasures.add("measureAccuracyTest");
        m_AdditionalMeasures.add("measureMeanAccuracyTest");
        m_AdditionalMeasures.add("measureAccuracyDiffTest");
        m_AdditionalMeasures.add("measureAUCTest");
        m_AdditionalMeasures.add("measureBrierScoreTest");
        m_AdditionalMeasures.add("measureBrierDiversityTest");
        measures = new DiversityMeasureImplementation().enumerateMeasures();
        while (measures.hasMoreElements()) {
            m_AdditionalMeasures.add(measures.nextElement() + "Test");
        }

        m_AdditionalMeasures.add("measureAccuracyOutOfBag");
        m_AdditionalMeasures.add("measureMeanAccuracyOutOfBag");
        m_AdditionalMeasures.add("measureAccuracyDiffOutOfBag");
        m_AdditionalMeasures.add("measureAUCOutOfBag");
        m_AdditionalMeasures.add("measureBrierScoreOutOfBag");
        m_AdditionalMeasures.add("measureBrierDiversityOutOfBag");
        measures = new DiversityMeasureImplementation().enumerateMeasures();
        while (measures.hasMoreElements()) {
            m_AdditionalMeasures.add(measures.nextElement() + "OutOfBag");
        }
        m_AdditionalMeasures.add("measureAccuracyTrainCl");
        m_AdditionalMeasures.add("measureMeanAccuracyTrainCl");
        m_AdditionalMeasures.add("measureAccuracyDiffTrainCl");
        m_AdditionalMeasures.add("measureAUCTrainCl");
        m_AdditionalMeasures.add("measureBrierScoreTrainCl");
        m_AdditionalMeasures.add("measureBrierDiversityTrainCl");
        measures = new DiversityMeasureImplementation().enumerateMeasures();
        while (measures.hasMoreElements()) {
            m_AdditionalMeasures.add(measures.nextElement() + "TrainCl");
        }

        m_AdditionalMeasures.add("measureAccuracyValidationCl");
        m_AdditionalMeasures.add("measureMeanAccuracyValidationCl");
        m_AdditionalMeasures.add("measureAccuracyDiffValidationCl");
        m_AdditionalMeasures.add("measureAUCValidationCl");
        m_AdditionalMeasures.add("measureBrierScoreValidationCl");
        m_AdditionalMeasures.add("measureBrierDiversityValidationCl");
        measures = new DiversityMeasureImplementation().enumerateMeasures();
        while (measures.hasMoreElements()) {
            m_AdditionalMeasures.add(measures.nextElement() + "ValidationCl");
        }

        m_AdditionalMeasures.add("measureAccuracyTestCl");
        m_AdditionalMeasures.add("measureMeanAccuracyTestCl");
        m_AdditionalMeasures.add("measureAccuracyDiffTestCl");
        m_AdditionalMeasures.add("measureAUCTestCl");
        m_AdditionalMeasures.add("measureBrierScoreTestCl");
        m_AdditionalMeasures.add("measureBrierDiversityTestCl");
        measures = new DiversityMeasureImplementation().enumerateMeasures();
        while (measures.hasMoreElements()) {
            m_AdditionalMeasures.add(measures.nextElement() + "TestCl");
        }

        return m_AdditionalMeasures.elements();
    }

    @Override
    public double getMeasure(String measureName) {
        if (measureName.matches("[\\w]*Train")) {
            String[] measure = measureName.split("Train");
            return getMeasure(measure[0], trainSet);
        }

        if (measureName.matches("[\\w]*Validation")) {
            String[] measure = measureName.split("Validation");
            return getMeasure(measure[0], validationSet);
        }

        if (measureName.matches("[\\w]*Test")) {
            String[] measure = measureName.split("Test");
            return getMeasure(measure[0], testSet);
        }

        if (m_Classifier instanceof IOutOfBagCalculator) {
            if (measureName.matches("[\\w]*OutOfBag")) {
                String[] measure = measureName.split("OutOfBag");
                return getMeasure(measure[0], outOfBagSet);
            }
        }
        if (measureName.matches("[\\w]*TrainCl")) {
            String[] measure = measureName.split("TrainCl");
            return getMeasure(measure[0], trainClassifierSet);
        }

        if (measureName.matches("[\\w]*ValidationCl")) {
            String[] measure = measureName.split("ValidationCl");
            return getMeasure(measure[0], validationClassifierSet);
        }

        if (measureName.matches("[\\w]*TestCl")) {
            String[] measure = measureName.split("TestCl");
            return getMeasure(measure[0], testClassifierSet);
        }
        return 0;
    }

    public double getMeasure(String measureName, SetEvaluator set) {
        if (measureName.matches("measureAccuracy")) {
            return set.getAccuracy();
        }

        if (measureName.matches("measureMeanAccuracy")) {
            return set.getBaseAccuracy();
        }

        if (measureName.matches("measureAUC")) {
            return Utils.sum(set.getAUC()) / set.getAUC().length;
        }

        if (measureName.matches("measureAccuracyDiff")) {
            return set.getAccuracy() - set.getBaseAccuracy();
        }

        if (measureName.matches("measureBrierScore")) {
            return set.getBrierScore();
        }

        if (measureName.matches("measureBrierDiversity")) {
            return set.getBrierDiversity();
        }

        return set.getDiversity().getMeasure(measureName);
    }

    public Enumeration listOptions() {

        Vector<Option> newVector = new Vector<Option>();

        newVector.addElement(new Option(
                "\tPercent of training set used as validation data", "V", 1,
                "-V"));

        Enumeration<Option> enu = (Enumeration) super.listOptions();
        while (enu.hasMoreElements()) {
            newVector.addElement(enu.nextElement());
        }
        return newVector.elements();
    }

    public void setOptions(String[] options) throws Exception {

        String str = Utils.getOption("V", options);
        if (str.length() != 0) {
            setValidationPercent(Double.parseDouble(str));
        } else {
            setValidationPercent(0);
        }

        super.setOptions(options);
    }

    /**
     * Gets the current settings of the Classifier.
     *
     * @return an array of strings suitable for passing to setOptions
     */
    public String[] getOptions() {

        String[] options;
        Vector<String> result = new Vector<String>();

        result.add("-V");
        result.add("" + getValidationPercent());

        options = super.getOptions();
        for (int i = 0; i < options.length; i++) {
            result.add(options[i]);
        }

        String[] array = new String[result.size()];
        for (int i = 0; i < result.size(); i++) {
            array[i] = result.elementAt(i).toString();
        }

        return array;
    }

    /**
     * Returns the tip text for this property
     *
     * @return tip text for this property suitable for displaying in the
     *         explorer/experimenter gui
     */
    public String validationPercentTipText() {
        return "The size of validation set, as a percentage of the training set size";
    }

    /**
     * Gets the size of validation set, as a percentage of the training set
     * size.
     *
     * @return the validation set size, as a percentage.
     */
    public double getValidationPercent() {

        return m_validationPercent;
    }

    /**
     * Sets the size of validation set, as a percentage of the training set
     * size.
     *
     * @param newValidationPercent
     *            the bag size, as a percentage.
     */
    public void setValidationPercent(double newValidationPercent) {

        m_validationPercent = newValidationPercent;
    }

    public static void main(String[] args) throws Exception {
        runClassifier(new ValidationSetEvaluator(), args);

    }
}
